#include "htf_ms_backend.h"
#include "../../utils/htf_crypto.h"
#include "../../utils/htf_logger.h"
#include "../../utils/htf_secure_mem.h"
#include <stdlib.h>
#include <string.h>
#include <stdio.h>

typedef struct htf_ms_backend {
    float* model_weights;
    size_t weight_count;
    int is_loaded;
    uint8_t* key;
    size_t key_len;
    size_t input_size;
    size_t output_size;
} htf_ms_backend_t;

static int ms_init(void* self, const char* model_path) {
    htf_ms_backend_t* b = (htf_ms_backend_t*)self;

    if (!b->key || b->key_len == 0) {
        HTF_LOG_ERROR("Key not set before init");
        return -1;
    }

    FILE* f = fopen(model_path, "rb");
    if (!f) return -1;
    fseek(f, 0, SEEK_END);
    size_t enc_size = ftell(f);
    fseek(f, 0, SEEK_SET);
    uint8_t* enc_buf = (uint8_t*)malloc(enc_size);
    if (!enc_buf) {
        fclose(f);
        return -1;
    }
    fread(enc_buf, 1, enc_size, f);
    fclose(f);

    uint8_t* plain = NULL;
    size_t plain_sz = 0;
    int ret = htf_crypto_decrypt(enc_buf, enc_size, b->key, b->key_len, &plain, &plain_sz);
    free(enc_buf);
    if (ret != 0) {
        return -2;
    }

    if (plain_sz % sizeof(float) != 0) {
        free(plain);
        return -3;
    }

    b->model_weights = (float*)plain;
    b->weight_count = plain_sz / sizeof(float);
    b->is_loaded = 1;
    HTF_LOG_INFO("MindSpore backend loaded %zu weights", b->weight_count);
    return 0;
}

static int ms_run(void* self, const void* input, void* output) {
    htf_ms_backend_t* b = (htf_ms_backend_t*)self;
    if (!b->is_loaded || !input || !output) return -1;

    const float* in = (const float*)input;
    float* out = (float*)output;

    for (size_t i = 0; i < b->output_size; i++) {
        float sum = 0.0f;
        for (size_t j = 0; j < b->input_size && j < b->weight_count; j++) {
            sum += in[j] * b->model_weights[(i * b->input_size + j) % b->weight_count];
        }
        out[i] = sum / 1000.0f;
    }

    htf_crypto_sha256(out, b->output_size * sizeof(float), NULL);
    return 0;
}

static void ms_destroy(void* self) {
    htf_ms_backend_t* b = (htf_ms_backend_t*)self;
    if (b) {
        if (b->model_weights) {
            htf_secure_memzero(b->model_weights, b->weight_count * sizeof(float));
            free(b->model_weights);
        }
        if (b->key) {
            htf_secure_memzero(b->key, b->key_len);
            free(b->key);
        }
        free(b);
    }
}

static const htf_ai_backend_ops_t ms_ops = {
    .init = ms_init,
    .run = ms_run,
    .destroy = ms_destroy
};

void* htf_ms_backend_create(void) {
    return calloc(1, sizeof(htf_ms_backend_t));
}

const htf_ai_backend_ops_t* htf_ms_backend_get_ops(void) {
    return &ms_ops;
}

int htf_ms_backend_set_config(void* impl, const uint8_t* key, size_t key_len,
                              size_t input_size, size_t output_size) {
    if (!impl) return -1;
    htf_ms_backend_t* b = (htf_ms_backend_t*)impl;

    if (key && key_len > 0) {
        b->key = (uint8_t*)malloc(key_len);
        if (!b->key) return -1;
        memcpy(b->key, key, key_len);
        b->key_len = key_len;
    }

    b->input_size = input_size;
    b->output_size = output_size;
    return 0;
}